package com.allwees.core.mq.rabbitmq.batch;

import cn.hutool.core.thread.ThreadUtil;
import com.rabbitmq.client.ShutdownSignalException;
import com.allwees.core.mq.rabbitmq.MqException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.listener.BlockingQueueConsumer;
import org.springframework.context.ApplicationEventPublisher;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 批量消息消费者（利用{@link BlockingQueueConsumer}）
 * <br/>子类需要继承{@link AbstractConsumer}，并实现方法{@link AbstractConsumer#consume(List)}
 *
 * @author yangylsky
 * @since 2019-11-25
 *
 */
@Slf4j
public class BatchMqConsumer<M, S> {
    /**
     * 批量处理消息Size
     */
	private int batchSize = 100;

	/**
	 * 批量处理消息超时时间，单位ms（如果队列未满但达到超时时间，也执行批处理）
	 */
	private long batchTimeout = 2000;

	private BlockingQueueConsumer blockingQueueConsumer;

	private AbstractConsumer<M, S> messageConsumer;

	private ExecutorService executor = Executors.newFixedThreadPool(1);

	private String modelName;

	private boolean started = false;

	BatchMqConsumer(int batchSize, long batchTimeout) {
		this(batchSize, batchTimeout, null);
	}

	BatchMqConsumer(int batchSize, long batchTimeout, AbstractConsumer<M, S> messageConsumer) {
		this.batchSize = batchSize;
		this.batchTimeout = batchTimeout;
		setMessageConsumer(messageConsumer);
	}

	void setBlockingQueueConsumer(BlockingQueueConsumer blockingQueueConsumer) {
		this.blockingQueueConsumer = blockingQueueConsumer;
	}

	public void setMessageConsumer(AbstractConsumer<M, S> messageConsumer) {
		if(started) {
			throw new MqException(this + " is started");
		}
		this.messageConsumer = messageConsumer;
		if(messageConsumer != null) {
			this.modelName = messageConsumer.mClass.getSimpleName();
		}
	}

	public void setEventPublisher(ApplicationEventPublisher applicationEventPublisher) {
		if(started) {
			throw new MqException(this + " is started");
		}
		blockingQueueConsumer.setApplicationEventPublisher(applicationEventPublisher);
	}

	public synchronized void start() {
		if(started) {
			throw new MqException(this + " is started");
		}
		if(messageConsumer == null) {
			throw new MqException("messageConsumer must has been set");
		}
		try {
			blockingQueueConsumer.start();

			executor.submit(new Callable<String>() {
				@Override
				public String call() throws Exception {
					execute();
					return null;
				}
			});
			started = true;
			log.info("{} is started", this);
		} catch(Exception e) {
			log.error(this + " start failed", e);
		}
	}

	private void execute() {
		List<S> list = new ArrayList<>();
		try {
			while(true) {
				int dealNum = 0;
				long batchStartTime = System.currentTimeMillis();
				try {
					int i = batchSize;
					while(i-- > 0) {
						Message msg = blockingQueueConsumer.nextMessage(batchTimeout);
						if(msg != null) {
							M m = messageConsumer.parseMessage(msg);
							list.add(messageConsumer.transform(m));
						}
						if(System.currentTimeMillis() - batchStartTime > batchTimeout) {
							break;
						}
					}
					if(!list.isEmpty()) {
						messageConsumer.consume(list);
						dealNum = list.size();
						list.clear();
						blockingQueueConsumer.commitIfNecessary(false);
					} else {
						//队列中没有拿到数据，sleep 1s
						ThreadUtil.safeSleep(1000);
					}
				} catch(IOException e) {
					log.error("", e);
				} catch(InterruptedException e) {
					log.error("", e);
					Thread.currentThread().interrupt();
				} finally {
					if(dealNum > 0) {
						messageConsumer.afterConsume(dealNum, System.currentTimeMillis() - batchStartTime);
					}
				}
			}
		} catch(ShutdownSignalException e) {
			log.error("", e);
			Thread.currentThread().interrupt();
		}
	}

	public void stop() {
		blockingQueueConsumer.stop();
	}

	@Override
	public String toString() {
		return String.format("BatchMessageConsumer[%s: %d]", modelName, batchSize);
	}
}
